"""随机划分训练集和验证集"""
import os
import random
from shutil import copyfile

# 输入路径
image_dir = r'G:\\yolov8\\data\\images'  # 替换成你的图像文件夹路径
label_dir = r'G:\\yolov8\\data\\label'  # 替换成你的标签文件夹路径

# 输出路径
train_image_dir = r'G:\\yolov8\\ultralytics-main\\ultralytics-main\\my_data\\detection\\train\\images'
train_label_dir = r'G:\\yolov8\\ultralytics-main\\ultralytics-main\\my_data\\detection\\train\\labels'
val_image_dir = r'G:\\yolov8\\ultralytics-main\\ultralytics-main\\my_data\\detection\\val\\images'
val_label_dir = r'G:\\yolov8\\ultralytics-main\\ultralytics-main\\my_data\\detection\\val\\labels'

# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)

# 划分数据集的比例
train_val_split = 0.8

# 获取图像文件列表
image_files = os.listdir(image_dir)
random.shuffle(image_files)

# 计算划分的索引
split_index = int(len(image_files) * train_val_split)

# 划分训练集和验证集
train_image_files = image_files[:split_index]
val_image_files = image_files[split_index:]

# 复制图像文件并相应地复制标签文件
def copy_images_and_labels(image_files, source_image_dir, source_label_dir, dest_image_dir, dest_label_dir):
    for image_file in image_files:
        # 复制图像文件
        source_image_path = os.path.join(source_image_dir, image_file)
        dest_image_path = os.path.join(dest_image_dir, image_file)
        copyfile(source_image_path, dest_image_path)

        # 复制对应的标签文件
        label_file = os.path.splitext(image_file)[0] + '.txt'
        source_label_path = os.path.join(source_label_dir, label_file)
        dest_label_path = os.path.join(dest_label_dir, label_file)
        copyfile(source_label_path, dest_label_path)

# 复制训练集图像和标签
copy_images_and_labels(train_image_files, image_dir, label_dir, train_image_dir, train_label_dir)

# 复制验证集图像和标签
copy_images_and_labels(val_image_files, image_dir, label_dir, val_image_dir, val_label_dir)
